import cv2
import os
import threading
from typing import Callable, Optional
import time
import keyboard


def extract_frames(rtsp_url, output_folder, frame_interval=1):
    """
    从RTSP流中捕获帧并保存到指定文件夹，每隔指定帧数抽取一帧。

    :param rtsp_url: RTSP流地址
    :param output_folder: 保存帧的文件夹路径
    :param frame_interval: 每隔几帧抽取一帧
    """
    # 打开RTSP流
    cap = cv2.VideoCapture(rtsp_url)

    # 检查RTSP流是否成功打开
    if not cap.isOpened():
        print("无法打开RTSP流")
        return

    # 创建保存帧的文件夹
    if not os.path.exists(output_folder):
        os.makedirs(output_folder)

    frame_count = 0
    saved_frame_count = 0

    # 读取帧并保存
    while True:
        ret, frame = cap.read()
        if not ret:
            break

        # 仅当当前帧计数是间隔值的倍数时保存帧
        if frame_count % frame_interval == 0:
            # 设置保存帧的文件名
            frame_filename = os.path.join(output_folder, f'frame_{saved_frame_count:05d}.jpg')

            # 保存帧
            cv2.imwrite(frame_filename, frame)

            saved_frame_count += 1

        frame_count += 1

    # 释放RTSP流
    cap.release()

    print("完成帧捕获和保存")


class RTSPFrameExtractor:
    """
    RTSP流帧提取器，支持在独立线程中运行和控制
    """
    def __init__(self, rtsp_url: str, frame_callback: Callable[[cv2.Mat, int], None]):
        """
        初始化RTSP帧提取器

        :param rtsp_url: RTSP流地址
        :param frame_callback: 帧数据回调函数，接收参数为(frame: cv2.Mat, frame_count: int)
        """
        self.rtsp_url = rtsp_url
        self.frame_callback = frame_callback
        self._thread: Optional[threading.Thread] = None
        self._running = False
        self._frame_count = 0
        self._cap: Optional[cv2.VideoCapture] = None

    def start(self) -> bool:
        """
        开始帧提取（在独立线程中运行）

        :return: 启动成功返回True，失败返回False
        """
        if self._running:
            print("帧提取已在运行中")
            return False

        self._running = True
        self._thread = threading.Thread(target=self._run, daemon=True)
        self._thread.start()
        return True

    def stop(self) -> None:
        """
        停止帧提取并释放资源
        """
        self._running = False
        if self._thread:
            self._thread.join(timeout=5.0)
            self._thread = None
        if self._cap:
            self._cap.release()
            self._cap = None
        print(f"帧提取已停止，共处理 {self._frame_count} 帧")
        self._frame_count = 0

    def _run(self) -> None:
        """
        帧提取线程主函数
        """
        # 打开RTSP流
        self._cap = cv2.VideoCapture(self.rtsp_url)

        # 检查RTSP流是否成功打开
        if not self._cap.isOpened():
            print("无法打开RTSP流")
            self._running = False
            return

        print(f"开始从RTSP流提取帧: {self.rtsp_url}")

        try:
            while self._running:
                ret, frame = self._cap.read()
                if not ret:
                    print("RTSP流读取失败或已结束")
                    break

                # 调用回调函数传递当前帧和帧计数
                self.frame_callback(frame, self._frame_count)
                self._frame_count += 1
                # break

        finally:
            # 确保资源释放
            if self._cap:
                self._cap.release()
                self._cap = None
            self._running = False



# 以下代码在示例使用部分添加按 esc 停止的功能
if __name__ == "__main__":

    def handle_frame(frame: cv2.Mat, frame_count: int):
        """帧处理回调函数示例，使用窗体展示图片"""
        print(f"处理第 {frame_count} 帧，尺寸: {frame.shape}")
        cv2.imshow('RTSP Frame', frame)
        cv2.waitKey(1)  # 等待1毫秒，使窗口能响应事件

    # 创建提取器实例
    extractor = RTSPFrameExtractor(
        rtsp_url="rtsp://admin:XG123456@192.168.1.154:554/streaming/channels/101",
        frame_callback=handle_frame
    )

    # 启动提取
    extractor.start()

    # 等待 esc 键按下停止提取
    while extractor._running:
        if keyboard.is_pressed('esc'):
            extractor.stop()
            break
        time.sleep(0.1)
     




# 仅当直接运行此脚本时执行以下代码
# if __name__ == "__main__":
#     rtsp_url = 'rtsp://admin:admin@10.10.18.171:8554/live/C123456C'
#     output_folder = 'frames'
#     frame_interval = 5  # 每隔5帧抽取一帧
#     extract_frames(rtsp_url, output_folder, frame_interval)
